{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-15T10:44:07.822873Z",
     "start_time": "2020-05-15T10:44:07.817744Z"
    }
   },
   "source": [
    "# ERNIE原理"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 对于中文，bert使用的基于字的处理，在mask时掩盖的也仅仅是一个单字。对于中文语料，学习到的模型能很简单地推测出字搭配，但是并不会学习到短语或者实体的语义信息。\n",
    "- ERNIE模型在BERT的基础上，加入了海量语料中的实体、短语等先验语义知识，建模真实世界的语义关系。在训练时将短语、实体等先验知识进行mask，强迫模型对其进行建模，学习它们的语义表示。               \n",
    "            \n",
    "    <img src=\"../images/ERNIE-mask.PNG\" width=\"80%\" alt=\"ERINE MASK\">\n",
    "       \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "具体来说， ERNIE 采用三种 masking 策略：\n",
    "\n",
    "    Basic-Level Masking： 跟 bert 一样对单字进行 mask，很难学习到高层次的语义信息；\n",
    "    Phrase-Level Masking： 输入仍然是单字级别的，mask连续短语；\n",
    "    Entity-Level Masking： 首先进行实体识别，然后将识别出的实体进行 mask。\n",
    "\n",
    "\n",
    "<img src=\"../images/ERNIE-masks.PNG\" width=\"100%\" alt=\"ERNIE masking\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "预训练阶段时采取上述掩码策略进行训练，获得每个汉字的嵌入向量；优调时，使用方法和BERT一样，**文本分词时仍然是按字分词**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-15T10:09:07.213301Z",
     "start_time": "2020-05-15T10:09:07.057141Z"
    }
   },
   "source": [
    "# 基于ERNIE模型的文本分类"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-15T12:28:14.683913Z",
     "start_time": "2020-05-15T12:28:13.563614Z"
    }
   },
   "outputs": [],
   "source": [
    "import time\n",
    "from tqdm import tqdm\n",
    "from datetime import timedelta\n",
    "\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from transformers import BertModel, BertTokenizer\n",
    "from transformers import AdamW, get_linear_schedule_with_warmup\n",
    "\n",
    "from sklearn import metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 模型配置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-15T12:28:19.100061Z",
     "start_time": "2020-05-15T12:28:19.092834Z"
    }
   },
   "outputs": [],
   "source": [
    "class Config(object):\n",
    "    \"\"\"配置参数\"\"\"\n",
    "    def __init__(self):\n",
    "        self.model_name = 'ERNIE'\n",
    "        dataset = \"../../H/datasets/THUCNews/\"\n",
    "        self.train_path = dataset + '/train.txt'  # 训练集\n",
    "        self.dev_path = dataset + '/dev.txt'  # 验证集\n",
    "        self.test_path = dataset + '/test.txt'  # 测试集\n",
    "        self.class_list = [\n",
    "            x.strip() for x in open(dataset + '/class.txt').readlines()\n",
    "        ]  # 类别名单\n",
    "        self.save_path = dataset + '/saved_dict/' + self.model_name + '.ckpt'  # 模型训练结果\n",
    "        self.device = torch.device(\n",
    "            'cuda' if torch.cuda.is_available() else 'cpu')  # 设备\n",
    "\n",
    "        self.require_improvement = 1000  # 若超过1000batch效果还没提升，则提前结束训练\n",
    "        self.num_classes = len(self.class_list)  # 类别数\n",
    "        self.num_epochs = 3  # epoch数\n",
    "        self.batch_size = 128  # mini-batch大小\n",
    "        self.pad_size = 32  # 每句话处理成的长度(短填长切)\n",
    "        self.learning_rate = 5e-5  # 学习率\n",
    "        self.bert_path = '../../H/models/huggingface/ERNIE/'\n",
    "        self.tokenizer = BertTokenizer.from_pretrained(self.bert_path)\n",
    "        self.hidden_size = 768"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-15T12:28:19.558703Z",
     "start_time": "2020-05-15T12:28:19.508824Z"
    }
   },
   "outputs": [],
   "source": [
    "config = Config()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 创建模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-15T12:28:21.663095Z",
     "start_time": "2020-05-15T12:28:21.656442Z"
    }
   },
   "outputs": [],
   "source": [
    "class Model(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super(Model, self).__init__()\n",
    "        self.bert = BertModel.from_pretrained(config.bert_path)\n",
    "        for param in self.bert.parameters():\n",
    "            param.requires_grad = True\n",
    "        self.fc = nn.Linear(config.hidden_size, config.num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        context = x[0]  # 输入的句子\n",
    "        mask = x[\n",
    "            2]  # 对padding部分进行mask，和句子一个size，padding部分用0表示，如：[1, 1, 1, 1, 0, 0]\n",
    "        _, pooled = self.bert(context,\n",
    "                              attention_mask=mask)\n",
    "        out = self.fc(pooled)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-15T12:28:23.138399Z",
     "start_time": "2020-05-15T12:28:21.967488Z"
    }
   },
   "outputs": [],
   "source": [
    "model = Model(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-15T12:28:23.158263Z",
     "start_time": "2020-05-15T12:28:23.152203Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Model(\n",
       "  (bert): BertModel(\n",
       "    (embeddings): BertEmbeddings(\n",
       "      (word_embeddings): Embedding(18000, 768, padding_idx=0)\n",
       "      (position_embeddings): Embedding(513, 768)\n",
       "      (token_type_embeddings): Embedding(2, 768)\n",
       "      (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (dropout): Dropout(p=0.1, inplace=False)\n",
       "    )\n",
       "    (encoder): BertEncoder(\n",
       "      (layer): ModuleList(\n",
       "        (0): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (1): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (2): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (3): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (4): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (5): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (6): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (7): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (8): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (9): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (10): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (11): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (pooler): BertPooler(\n",
       "      (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "      (activation): Tanh()\n",
       "    )\n",
       "  )\n",
       "  (fc): Linear(in_features=768, out_features=10, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 创建数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-15T12:28:30.379541Z",
     "start_time": "2020-05-15T12:28:30.360648Z"
    }
   },
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "def build_dataset(config):\n",
    "    def load_dataset(path, pad_size=32):\n",
    "        contents = []\n",
    "        with open(path, 'r', encoding='UTF-8') as f:\n",
    "            for line in tqdm(f):\n",
    "                lin = line.strip()\n",
    "                if not lin:\n",
    "                    continue\n",
    "                content, label = lin.split('\\t')\n",
    "\n",
    "                token_ids = config.tokenizer.encode(\n",
    "                    content,\n",
    "                    add_special_tokens=True,  # 添加特殊符号\n",
    "                )\n",
    "\n",
    "                mask = []\n",
    "                seq_len = len(token_ids)\n",
    "\n",
    "                if pad_size:\n",
    "                    if seq_len < pad_size:\n",
    "                        mask = [1] * len(token_ids) + [0\n",
    "                                                       ] * (pad_size - seq_len)\n",
    "                        token_ids += ([0] * (pad_size - seq_len))\n",
    "                    else:\n",
    "                        mask = [1] * pad_size\n",
    "                        token_ids = token_ids[:pad_size]\n",
    "                        seq_len = pad_size\n",
    "                contents.append((token_ids, int(label), seq_len, mask))\n",
    "        return contents\n",
    "\n",
    "    train = load_dataset(config.train_path, config.pad_size)\n",
    "    dev = load_dataset(config.dev_path, config.pad_size)\n",
    "    test = load_dataset(config.test_path, config.pad_size)\n",
    "    return train, dev, test\n",
    "\n",
    "\n",
    "class DatasetIterater(object):\n",
    "    def __init__(self, batches, batch_size, device):\n",
    "        self.batch_size = batch_size\n",
    "        self.batches = batches\n",
    "        self.n_batches = len(batches) // batch_size\n",
    "        self.residue = False  # 记录batch数量是否为整数\n",
    "        if len(batches) % self.n_batches != 0:\n",
    "            self.residue = True\n",
    "        self.index = 0\n",
    "        self.device = device\n",
    "\n",
    "    def _to_tensor(self, datas):\n",
    "        x = torch.LongTensor([_[0] for _ in datas]).to(self.device)\n",
    "        y = torch.LongTensor([_[1] for _ in datas]).to(self.device)\n",
    "\n",
    "        # pad前的长度(超过pad_size的设为pad_size)\n",
    "        seq_len = torch.LongTensor([_[2] for _ in datas]).to(self.device)\n",
    "        mask = torch.LongTensor([_[3] for _ in datas]).to(self.device)\n",
    "        return (x, seq_len, mask), y\n",
    "\n",
    "    def __next__(self):\n",
    "        if self.residue and self.index == self.n_batches:\n",
    "            batches = self.batches[self.index *\n",
    "                                   self.batch_size:len(self.batches)]\n",
    "            self.index += 1\n",
    "            batches = self._to_tensor(batches)\n",
    "            return batches\n",
    "\n",
    "        elif self.index >= self.n_batches:\n",
    "            self.index = 0\n",
    "            raise StopIteration\n",
    "        else:\n",
    "            batches = self.batches[self.index *\n",
    "                                   self.batch_size:(self.index + 1) *\n",
    "                                   self.batch_size]\n",
    "            self.index += 1\n",
    "            batches = self._to_tensor(batches)\n",
    "            return batches\n",
    "\n",
    "    def __iter__(self):\n",
    "        return self\n",
    "\n",
    "    def __len__(self):\n",
    "        if self.residue:\n",
    "            return self.n_batches + 1\n",
    "        else:\n",
    "            return self.n_batches\n",
    "\n",
    "\n",
    "def build_iterator(dataset, config):\n",
    "    iter = DatasetIterater(dataset, config.batch_size, config.device)\n",
    "    return iter\n",
    "\n",
    "\n",
    "def get_time_dif(start_time):\n",
    "    \"\"\"获取已使用时间\"\"\"\n",
    "    end_time = time.time()\n",
    "    time_dif = end_time - start_time\n",
    "    return timedelta(seconds=int(round(time_dif)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-15T12:28:50.474281Z",
     "start_time": "2020-05-15T12:28:30.668189Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "180000it [00:17, 10078.27it/s]\n",
      "10000it [00:00, 10299.09it/s]\n",
      "10000it [00:00, 10382.67it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 19.8 s, sys: 74.2 ms, total: 19.9 s\n",
      "Wall time: 19.8 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "train_data, dev_data, test_data = build_dataset(config)\n",
    "train_iter = build_iterator(train_data, config)\n",
    "dev_iter = build_iterator(dev_data, config)\n",
    "test_iter = build_iterator(test_data, config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 训练模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-15T12:29:50.143769Z",
     "start_time": "2020-05-15T12:29:50.133387Z"
    }
   },
   "outputs": [],
   "source": [
    "def train(config, model, train_iter, dev_iter, test_iter):\n",
    "    start = time.time()\n",
    "    \n",
    "    # 优化器\n",
    "    optimizer = AdamW(model.parameters(), \n",
    "                      lr=config.learning_rate, \n",
    "                      correct_bias=False)  # To reproduce BertAdam specific behavior set correct_bias=False\n",
    "    total_steps = len(train_iter) * config.num_epochs\n",
    "    scheduler = get_linear_schedule_with_warmup(optimizer, \n",
    "                                                num_warmup_steps=0, \n",
    "                                                num_training_steps=total_steps)  # PyTorch scheduler\n",
    "\n",
    "    total_batch = 0  # 记录进行到多少batch\n",
    "    dev_best_loss = float('inf')\n",
    "    last_improve = 0  # 记录上次验证集loss下降的batch数\n",
    "    flag = False  # 记录是否很久没有效果提升\n",
    "\n",
    "    # 训练模式\n",
    "    model.train()\n",
    "    for epoch in range(config.num_epochs):\n",
    "        print('Epoch [{}/{}]'.format(epoch + 1, config.num_epochs))\n",
    "\n",
    "        for i, (trains, labels) in enumerate(train_iter):\n",
    "            # 前向推理\n",
    "            outputs = model(trains)\n",
    "\n",
    "            # 梯度归零\n",
    "            model.zero_grad()\n",
    "\n",
    "            # 损失\n",
    "            loss = F.cross_entropy(outputs, labels)\n",
    "\n",
    "            # 反向传播\n",
    "            loss.backward()\n",
    "            \n",
    "            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
    "\n",
    "            # 更新参数\n",
    "            optimizer.step()\n",
    "            scheduler.step()\n",
    "\n",
    "            # 精度\n",
    "            if total_batch % 100 == 0:\n",
    "                # 每多少轮输出在训练集和验证集上的效果\n",
    "                true = labels.data.cpu()\n",
    "                predic = torch.max(outputs.data, 1)[1].cpu()\n",
    "                train_acc = metrics.accuracy_score(true, predic)\n",
    "\n",
    "                # 验证集上的精度和损失\n",
    "                dev_acc, dev_loss = evaluate(config, model, dev_iter)\n",
    "                if dev_loss < dev_best_loss:\n",
    "                    dev_best_loss = dev_loss\n",
    "\n",
    "                    # 保存最佳模型\n",
    "                    torch.save(model.state_dict(), config.save_path)\n",
    "\n",
    "                    improve = '*'\n",
    "                    last_improve = total_batch\n",
    "                else:\n",
    "                    improve = ''\n",
    "\n",
    "                time_dif = get_time_dif(start)\n",
    "\n",
    "                msg = 'Iter: {0:>6},  Train Loss: {1:>5.2},  Train Acc: {2:>6.2%},  Val Loss: {3:>5.2},  Val Acc: {4:>6.2%},  Time: {5} {6}'\n",
    "                print(\n",
    "                    msg.format(total_batch, loss.item(), train_acc, dev_loss,\n",
    "                               dev_acc, time_dif, improve))\n",
    "\n",
    "                model.train()\n",
    "\n",
    "            total_batch += 1\n",
    "            if total_batch - last_improve > config.require_improvement:\n",
    "                # 验证集loss超过1000batch没下降，结束训练\n",
    "                print(\"No optimization for a long time, auto-stopping...\")\n",
    "                flag = True\n",
    "                break\n",
    "        if flag:\n",
    "            break\n",
    "\n",
    "    # 测试模型\n",
    "    test(config, model, test_iter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-15T12:29:52.144913Z",
     "start_time": "2020-05-15T12:29:52.139209Z"
    }
   },
   "outputs": [],
   "source": [
    "# 测试模型\n",
    "\n",
    "\n",
    "def test(config, model, test_iter):\n",
    "    # 加载模型\n",
    "    model.load_state_dict(torch.load(config.save_path))\n",
    "\n",
    "    # 验证模式\n",
    "    model.eval()\n",
    "\n",
    "    start_time = time.time()\n",
    "\n",
    "    test_acc, test_loss, test_report, test_confusion = evaluate(config,\n",
    "                                                                model,\n",
    "                                                                test_iter,\n",
    "                                                                test=True)\n",
    "    msg = 'Test Loss: {0:>5.2},  Test Acc: {1:>6.2%}'\n",
    "    print(msg.format(test_loss, test_acc))\n",
    "    print(\"Precision, Recall and F1-Score...\")\n",
    "    print(test_report)\n",
    "    print(\"Confusion Matrix...\")\n",
    "    print(test_confusion)\n",
    "    time_dif = get_time_dif(start_time)\n",
    "    print(\"Time usage:\", time_dif)\n",
    "\n",
    "\n",
    "    # 验证模型\n",
    "def evaluate(config, model, data_iter, test=False):\n",
    "    model.eval()\n",
    "    loss_total = 0\n",
    "    predict_all = np.array([], dtype=int)\n",
    "    labels_all = np.array([], dtype=int)\n",
    "    with torch.no_grad():\n",
    "        for texts, labels in data_iter:\n",
    "            outputs = model(texts)\n",
    "            loss = F.cross_entropy(outputs, labels)\n",
    "            loss_total += loss\n",
    "            labels = labels.data.cpu().numpy()\n",
    "            predic = torch.max(outputs.data, 1)[1].cpu().numpy()\n",
    "            labels_all = np.append(labels_all, labels)\n",
    "            predict_all = np.append(predict_all, predic)\n",
    "\n",
    "    acc = metrics.accuracy_score(labels_all, predict_all)\n",
    "    if test:\n",
    "        report = metrics.classification_report(labels_all,\n",
    "                                               predict_all,\n",
    "                                               target_names=config.class_list,\n",
    "                                               digits=4)\n",
    "        confusion = metrics.confusion_matrix(labels_all, predict_all)\n",
    "        return acc, loss_total / len(data_iter), report, confusion\n",
    "    return acc, loss_total / len(data_iter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-15T12:42:07.368522Z",
     "start_time": "2020-05-15T12:29:54.133822Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/3]\n",
      "Iter:      0,  Train Loss:   2.4,  Train Acc:  5.47%,  Val Loss:   2.3,  Val Acc: 22.21%,  Time: 0:00:07 *\n",
      "Iter:    100,  Train Loss:   0.3,  Train Acc: 92.19%,  Val Loss:  0.34,  Val Acc: 90.17%,  Time: 0:00:36 *\n",
      "Iter:    200,  Train Loss:   0.3,  Train Acc: 89.84%,  Val Loss:  0.27,  Val Acc: 91.57%,  Time: 0:01:06 *\n",
      "Iter:    300,  Train Loss:  0.23,  Train Acc: 92.19%,  Val Loss:  0.27,  Val Acc: 91.61%,  Time: 0:01:36 \n",
      "Iter:    400,  Train Loss:  0.35,  Train Acc: 89.06%,  Val Loss:  0.24,  Val Acc: 92.15%,  Time: 0:02:06 *\n",
      "Iter:    500,  Train Loss:  0.16,  Train Acc: 94.53%,  Val Loss:  0.23,  Val Acc: 92.63%,  Time: 0:02:36 *\n",
      "Iter:    600,  Train Loss:  0.29,  Train Acc: 90.62%,  Val Loss:  0.22,  Val Acc: 92.83%,  Time: 0:03:07 *\n",
      "Iter:    700,  Train Loss:  0.14,  Train Acc: 94.53%,  Val Loss:  0.21,  Val Acc: 93.00%,  Time: 0:03:37 *\n",
      "Iter:    800,  Train Loss:  0.16,  Train Acc: 95.31%,  Val Loss:  0.21,  Val Acc: 93.12%,  Time: 0:04:07 *\n",
      "Iter:    900,  Train Loss:  0.21,  Train Acc: 91.41%,  Val Loss:   0.2,  Val Acc: 93.46%,  Time: 0:04:37 *\n",
      "Iter:   1000,  Train Loss:  0.12,  Train Acc: 95.31%,  Val Loss:  0.21,  Val Acc: 93.23%,  Time: 0:05:07 \n",
      "Iter:   1100,  Train Loss:  0.22,  Train Acc: 92.97%,  Val Loss:  0.19,  Val Acc: 93.68%,  Time: 0:05:37 *\n",
      "Iter:   1200,  Train Loss:   0.2,  Train Acc: 92.97%,  Val Loss:  0.19,  Val Acc: 93.60%,  Time: 0:06:08 *\n",
      "Iter:   1300,  Train Loss:  0.22,  Train Acc: 89.84%,  Val Loss:  0.18,  Val Acc: 93.94%,  Time: 0:06:38 *\n",
      "Iter:   1400,  Train Loss:  0.35,  Train Acc: 92.19%,  Val Loss:  0.18,  Val Acc: 94.13%,  Time: 0:07:08 *\n",
      "Epoch [2/3]\n",
      "Iter:   1500,  Train Loss:   0.1,  Train Acc: 96.09%,  Val Loss:  0.19,  Val Acc: 94.12%,  Time: 0:07:38 \n",
      "Iter:   1600,  Train Loss:  0.13,  Train Acc: 93.75%,  Val Loss:  0.19,  Val Acc: 94.15%,  Time: 0:08:08 \n",
      "Iter:   1700,  Train Loss:  0.14,  Train Acc: 95.31%,  Val Loss:  0.19,  Val Acc: 93.86%,  Time: 0:08:38 \n",
      "Iter:   1800,  Train Loss:  0.13,  Train Acc: 95.31%,  Val Loss:  0.19,  Val Acc: 93.90%,  Time: 0:09:07 \n",
      "Iter:   1900,  Train Loss: 0.099,  Train Acc: 94.53%,  Val Loss:  0.19,  Val Acc: 93.99%,  Time: 0:09:37 \n",
      "Iter:   2000,  Train Loss:  0.15,  Train Acc: 95.31%,  Val Loss:  0.19,  Val Acc: 94.08%,  Time: 0:10:07 \n",
      "Iter:   2100,  Train Loss:  0.08,  Train Acc: 97.66%,  Val Loss:  0.19,  Val Acc: 94.10%,  Time: 0:10:37 \n",
      "Iter:   2200,  Train Loss:  0.15,  Train Acc: 93.75%,  Val Loss:  0.19,  Val Acc: 94.13%,  Time: 0:11:07 \n",
      "Iter:   2300,  Train Loss: 0.059,  Train Acc: 97.66%,  Val Loss:  0.19,  Val Acc: 94.10%,  Time: 0:11:37 \n",
      "Iter:   2400,  Train Loss: 0.086,  Train Acc: 96.88%,  Val Loss:  0.19,  Val Acc: 94.16%,  Time: 0:12:07 \n",
      "No optimization for a long time, auto-stopping...\n",
      "Test Loss:  0.17,  Test Acc: 94.44%\n",
      "Precision, Recall and F1-Score...\n",
      "               precision    recall  f1-score   support\n",
      "\n",
      "      finance     0.9376    0.9320    0.9348      1000\n",
      "       realty     0.9666    0.9550    0.9608      1000\n",
      "       stocks     0.9040    0.9040    0.9040      1000\n",
      "    education     0.9634    0.9730    0.9682      1000\n",
      "      science     0.9131    0.9140    0.9135      1000\n",
      "      society     0.9538    0.9300    0.9418      1000\n",
      "     politics     0.9136    0.9310    0.9222      1000\n",
      "       sports     0.9830    0.9840    0.9835      1000\n",
      "         game     0.9563    0.9620    0.9591      1000\n",
      "entertainment     0.9533    0.9590    0.9561      1000\n",
      "\n",
      "     accuracy                         0.9444     10000\n",
      "    macro avg     0.9445    0.9444    0.9444     10000\n",
      " weighted avg     0.9445    0.9444    0.9444     10000\n",
      "\n",
      "Confusion Matrix...\n",
      "[[932   9  39   2   4   1   7   2   4   0]\n",
      " [ 10 955   9   2   3   5   7   2   1   6]\n",
      " [ 35  11 904   0  17   2  27   2   1   1]\n",
      " [  2   1   1 973   2   9   4   1   0   7]\n",
      " [  2   1  21   5 914  11   9   1  24  12]\n",
      " [  9   7   1  11   4 930  25   0   3  10]\n",
      " [  3   1  23   9  16  12 931   1   1   3]\n",
      " [  0   1   0   2   1   1   5 984   2   4]\n",
      " [  0   0   1   0  27   3   2   1 962   4]\n",
      " [  1   2   1   6  13   1   2   7   8 959]]\n",
      "Time usage: 0:00:06\n"
     ]
    }
   ],
   "source": [
    "model = model.to(config.device)\n",
    "train(config, model, train_iter, dev_iter, test_iter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
